import gradio as gr
import sqlite3
import bcrypt
import jwt
import torch
import datetime
from transformers import GenerationConfig
from qwen.modeling_qwen import QWenLMHeadModel
from qwen.tokenization_qwen import QWenTokenizer
from arguments import GenegrateArguments


# 初始化模型
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

genegrate_args = GenegrateArguments()

model = QWenLMHeadModel.from_pretrained(genegrate_args.model_dir).to(device)

tokenizer = QWenTokenizer.from_pretrained(genegrate_args.tokenizer_dir)

gen_config = GenerationConfig(
    temperature=genegrate_args.temperature,
    top_k=genegrate_args.top_k,
    top_p=genegrate_args.top_p,
    do_sample=genegrate_args.do_sample,
    num_beams=genegrate_args.num_beams,
    repetition_penalty=genegrate_args.repetition_penalty,
    max_new_tokens=genegrate_args.max_new_tokens,
    eos_token_id=tokenizer.eos_token_id,
    pad_token_id=tokenizer.pad_token_id,
)

model.eval()

state = {"token": None, "username": None}

# JWT 密钥与有效期
SECRET_KEY = "your_secret_key"
TOKEN_EXPIRY_HOURS = 2

# 初始化数据库
def init_db():
    conn = sqlite3.connect("users.db")
    cursor = conn.cursor()

    # 创建用户表
    cursor.execute('''CREATE TABLE IF NOT EXISTS users (
                      id INTEGER PRIMARY KEY AUTOINCREMENT,
                      username TEXT UNIQUE,
                      password TEXT)''')

    # 创建对话记录表
    cursor.execute('''CREATE TABLE IF NOT EXISTS conversations (
                      id INTEGER PRIMARY KEY AUTOINCREMENT,
                      user_id INTEGER,
                      user_input TEXT,
                      model_response TEXT,
                      timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
                      FOREIGN KEY (user_id) REFERENCES users(id))''')

    conn.commit()
    conn.close()

# 生成 JWT 令牌
def generate_jwt(username):
    payload = {
        "username": username,
        "exp": datetime.datetime.utcnow() + datetime.timedelta(hours=TOKEN_EXPIRY_HOURS)
    }
    return jwt.encode(payload, SECRET_KEY, algorithm="HS256")

# 验证 JWT 令牌
def verify_jwt(token):
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=["HS256"])
        return payload["username"]
    except (jwt.ExpiredSignatureError, jwt.InvalidTokenError):
        return None

# 用户注册
def register(username, password):
    conn = sqlite3.connect("users.db")
    cursor = conn.cursor()
    try:
        hashed_pw = bcrypt.hashpw(password.encode('utf-8'), bcrypt.gensalt())
        cursor.execute("INSERT INTO users (username, password) VALUES (?, ?)", (username, hashed_pw))
        conn.commit()
        return "✅ 注册成功，请登录！"
    except sqlite3.IntegrityError:
        return "❌ 用户名已存在，请更换！"
    finally:
        conn.close()

# 用户登录
def login(username, password):
    conn = sqlite3.connect("users.db")
    cursor = conn.cursor()
    cursor.execute("SELECT id, password FROM users WHERE username=?", (username,))
    user = cursor.fetchone()
    conn.close()

    if user and bcrypt.checkpw(password.encode('utf-8'), user[1]):
        token = generate_jwt(username)
        state["token"] = token
        state["username"] = username
        return f"🎉 欢迎 {username}，登录成功！"
    else:
        return "❌ 登录失败，用户名或密码错误！"

# 用户退出登录
def logout():
    state["token"] = None
    state["username"] = None
    return gr.update(visible=True), gr.update(visible=False)

# 记录对话数据
def record_conversation(user_id, user_input, model_response):
    conn = sqlite3.connect("users.db")
    cursor = conn.cursor()
    cursor.execute("INSERT INTO conversations (user_id, user_input, model_response) VALUES (?, ?, ?)", 
                   (user_id, user_input, model_response))
    conn.commit()
    conn.close()

# 处理对话
def chat_with_model(user_input, chat_history):
    username = verify_jwt(state['token'])
    if not username:
        return "❌ 未授权访问，请重新登录！", chat_history

    # 获取用户 ID
    conn = sqlite3.connect("users.db")
    cursor = conn.cursor()
    cursor.execute("SELECT id FROM users WHERE username=?", (username,))
    user = cursor.fetchone()
    user_id = user[0]
    conn.close()

    text = "你是一个助手 用户: {text} 回答: ".format(text=user_input)
    tokend = tokenizer(text, add_special=False)
    input_ids, attention_mask = torch.LongTensor([tokend.input_ids]).to(device), torch.LongTensor([tokend.attention_mask]).to(device)
    outputs = model.generate(inputs=input_ids, attention_mask=attention_mask, generation_config=gen_config)
    outs = tokenizer.decode(outputs[0].cpu().numpy())
    model_response = outs.split("<|im_end|>")[0].split("回答: ")[1]

    # 记录对话
    record_conversation(user_id, user_input, model_response)

    # 更新对话历史
    chat_history.append(("User", user_input))
    chat_history.append(("Model", model_response))

    return model_response

# 查看用户历史对话记录
def get_conversation_history():
    username = verify_jwt(state["token"])
    if not username:
        return "❌ 未授权访问，请重新登录！"

    conn = sqlite3.connect("users.db")
    cursor = conn.cursor()
    cursor.execute("SELECT id FROM users WHERE username=?", (username,))
    user = cursor.fetchone()
    user_id = user[0]

    cursor.execute("SELECT user_input, model_response, timestamp FROM conversations WHERE user_id=? ORDER BY timestamp DESC", (user_id,))
    conversations = cursor.fetchall()
    conn.close()

    if conversations:
        return "\n\n".join([f"时间: {conv[2]}\n用户: {conv[0]}\n模型: {conv[1]}" for conv in conversations])
    else:
        return "❌ 没有历史对话记录。"

# 初始化数据库
init_db()

# 构建 Gradio 界面
with gr.Blocks() as app:
    gr.Markdown("# 🔐 用户认证 & 大模型对话")

    # 注册界面
    with gr.Tab("注册"):
        reg_username = gr.Textbox(label="用户名")
        reg_password = gr.Textbox(label="密码", type="password")
        reg_button = gr.Button("注册")
        reg_output = gr.Textbox(label="注册结果")
        reg_button.click(register, inputs=[reg_username, reg_password], outputs=reg_output)

    # 登录界面
    with gr.Tab("登录"):
        login_username = gr.Textbox(label="用户名")
        login_password = gr.Textbox(label="密码", type="password")
        login_button = gr.Button("登录")
        login_output = gr.Textbox(label="登录反馈")
        login_button.click(login, inputs=[login_username, login_password], outputs=login_output)

    # 大模型对话页面（使用 ChatInterface，需登录）
    with gr.Tab("大模型对话"):
        token_input_chat = gr.Textbox(label="JWT 令牌（请先登录获取）", visible=False)
        gr.ChatInterface(fn=chat_with_model, 
                         chatbot=gr.Chatbot(height=500, value=[['你好', '你好，我是你的助手，我将尽力帮助你解决问题']]),
                         textbox=gr.Textbox(placeholder='请输入你的问题', container=False),
                         title='AI助手',
                         theme='soft',
                         examples=['你是谁', '使用Java实现快速排序'],
                         retry_btn=None,
                         submit_btn='发送',
                         undo_btn='删除前言',
                         clear_btn='清空',
                         ).queue()
        logout_button = gr.Button("退出登录")
        logout_button.click(logout, outputs=[login_output, token_input_chat])

    # 查看历史对话
    with gr.Tab("历史对话"):
        history_button = gr.Button("查看历史对话")
        history_output = gr.Textbox(label="历史对话记录")
        history_button.click(get_conversation_history, outputs=history_output)

app.launch()